# -*- coding: utf-8 -*
from flyai.framework import FlyAI
import cv2
import numpy as np

from path import *
from model import get_models, get_model_by_name
from config import args, cfg

import torch
from torchvision import transforms


class Prediction(FlyAI):
    def __init__(self, arguments=args):
        self.args = arguments
        self.h, self.w = self.args.height, self.args.width
        self.mean, self.std = np.load(MEAN_PATH), np.load(STD_PATH)
        # self.mean, self.std = [0.3876, 0.4219, 0.4984], [0.2931, 0.2922, 0.3144]  # BGR
        self.transform = transforms.Compose([transforms.ToTensor(),
                                             transforms.Normalize(mean=self.mean, std=self.std)])

    def load_model(self):
        n_models = args.n_models

        # save path
        if args.save_mode == 'separate':
            save_path = [eval(f'MODEL_FILE_PATH{i + 1}2') for i in range(n_models)]
        elif args.save_mode == 'ensemble':
            save_path = [eval(f'MODEL_FILE_PATH{i + 1}1') for i in range(n_models)]

        for i in range(n_models):
            exec(f'self.model_{i + 1} = get_model_by_name(cfg.model{i + 1}, pretrain = False)')
            exec(f'self.model_{i + 1}.load_state_dict(torch.load(save_path[{i}]))')
            exec(f'self.model_{i + 1}.eval()')

    def add(self, base: str):

        return eval(' + '.join([f'{base}_{i + 1}' for i in range(args.n_models)]))

    def predict(self, image_path):

        if type(image_path) is dict:
            image_path = image_path["image_path"]

        img = cv2.imread(image_path)
        img = cv2.resize(img, (self.w, self.h), interpolation=cv2.INTER_LINEAR)
        img = self.transform(img).unsqueeze(0).cuda()

        with torch.no_grad():
            for i in range(args.n_models):
                exec(f'self.prediction_{i + 1} = self.model_{i + 1}(img)')
            prediction_ens = self.add('self.prediction')

            label = int(prediction_ens.argmax(1))

        return {"label": label}


if __name__ == '__main__':
    predict = Prediction()
    predict.load_model()
    r = predict.predict(image_path=r'F:\FlyAI\BaldClassification_FlyAI\data\input\BaldClassification\image\28.jpg')
    print(r)
